import asyncio
import logging
from typing import Dict, List, Any, Optional

from mcpTool.MCPClient import MCPClient

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("mcp_manager")


class MCPClientManager:
    """MCP客户端管理器（单例模式），维护一个长期运行的MCP客户端实例"""
    
    _instance = None
    _lock = asyncio.Lock()
    _initialized = False
    
    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super(MCPClientManager, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self, client_name: str = "Wend MCP Client", client_version: str = "1.0.0"):
        """初始化MCP客户端管理器基础属性（非异步操作）
        
        Args:
            client_name: 客户端名称
            client_version: 客户端版本
        """
        # 只设置基本属性，不执行异步操作
        self.client = None
        self.current_connection = None
        self.is_connected = None
        self.client_name = client_name
        self.client_version = client_version
        # 其他属性将在initialize()中初始化
    
    async def initialize(self):
        """异步初始化MCP客户端管理器
        
        必须在使用管理器前调用此方法
        
        Returns:
            self: 返回自身实例以支持链式调用
        """
        async with self._lock:
            if self._initialized:
                return self
                
            self.client = MCPClient(self.client_name, self.client_version)
            self.current_connection = None  # 当前连接的服务器信息
            self.is_connected = False
            self._initialized = True
            logger.info(f"创建MCP客户端管理器: {self.client_name} v{self.client_version}")
            return self
    
    async def connect_sse(self, base_url: str, session_id: Optional[str] = None):
        """使用SSE连接到MCP服务器
        
        Args:
            base_url: MCP服务器基础URL，例如 "http://localhost:3000"
            session_id: 会话ID，用于区分不同的客户端连接
            
        Returns:
            连接成功返回True，否则返回False
        """
        if not self._initialized:
            await self.initialize()
            
        async with self._lock:
            # 如果已连接，先断开现有连接
            if self.is_connected:
                await self.disconnect()
                
            try:
                # 确保 base_url 格式正确
                if not base_url.startswith(('http://', 'https://')):
                    base_url = 'https://' + base_url

                # 尝试连接，最多重试3次
                max_retries = 3
                retry_count = 0
                last_error = None

                while retry_count < max_retries:
                    try:
                        await self.client.connect_with_sse(base_url, session_id)
                        self.current_connection = {
                            "connection_type": "sse",
                            "base_url": base_url,
                            "session_id": session_id
                        }
                        self.is_connected = True
                        logger.info(f"已使用SSE连接到MCP服务器: {base_url}")
                        return True
                    except Exception as e:
                        last_error = e
                        retry_count += 1
                        if retry_count < max_retries:
                            logger.warning(f"SSE连接失败 (尝试 {retry_count}/{max_retries}): {str(e)}")
                            await asyncio.sleep(1)  # 等待1秒后重试
                        else:
                            logger.error(f"SSE连接MCP服务器失败: {str(last_error)}")
                            return False
            except Exception as e:
                logger.error(f"SSE连接MCP服务器失败: {str(e)}")
                return False
    
    async def connect(self, command: str, args: Optional[List[str]] = None, env: Optional[Dict[str, str]] = None):
        """连接到MCP服务器
        
        Args:
            command: 命令
            args: 命令参数列表
            env: 环境变量
            
        Returns:
            连接成功返回True，否则返回False
        """
        if not self._initialized:
            await self.initialize()
            
        async with self._lock:
            # 如果已连接，先断开现有连接
            if self.is_connected:
                await self.disconnect()
                
            try:
                await self.client.connect_with_command(command, args, env)
                self.current_connection = {
                    "connection_type": "stdio",
                    "command": command,
                    "args": args,
                    "env": env
                }
                self.is_connected = True
                logger.info(f"已连接到MCP服务器: {command} {' '.join(args or [])}")
                return True
            except Exception as e:
                print(e)
                logger.error(f"连接MCP服务器失败: {str(e)}")
                return False
    
    async def connect_script(self, script_path: str):
        """连接到基于脚本的MCP服务器
        
        Args:
            script_path: 脚本路径
            
        Returns:
            连接成功返回True，否则返回False
        """
        if not self._initialized:
            await self.initialize()
            
        async with self._lock:
            # 如果已连接，先断开现有连接
            if self.is_connected:
                await self.disconnect()
                
            try:
                await self.client.connect_with_script(script_path)
                self.current_connection = {
                    "connection_type": "stdio",
                    "script_path": script_path
                }
                self.is_connected = True
                logger.info(f"已连接到脚本MCP服务器: {script_path}")
                return True
            except Exception as e:
                logger.error(f"连接脚本MCP服务器失败: {str(e)}")
                return False
    
    async def disconnect(self):
        """断开当前连接"""
        if not self._initialized:
            await self.initialize()
            
        async with self._lock:
            if self.is_connected:
                await self.client.close()
                self.is_connected = False
                logger.info("已断开MCP服务器连接")
    
    async def reconnect(self):
        """重新连接到上一个服务器"""
        if not self._initialized:
            await self.initialize()
            
        async with self._lock:
            if not self.current_connection:
                logger.error("没有可重连的服务器")
                return False
            
            connection_type = self.current_connection.get("connection_type", "stdio")
            
            if connection_type == "sse":
                return await self.connect_sse(
                    self.current_connection["base_url"],
                    self.current_connection.get("session_id")
                )
            elif "script_path" in self.current_connection:
                return await self.connect_script(self.current_connection["script_path"])
            else:
                return await self.connect(
                    self.current_connection["command"],
                    self.current_connection["args"],
                    self.current_connection["env"]
                )
    
    async def call_tool(self, name: str, arguments: Dict[str, Any]):
        """调用MCP工具
        
        Args:
            name: 工具名称
            arguments: 工具参数
            
        Returns:
            工具调用结果
        """
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect、connect_script或connect_sse方法")
        
        return await self.client.call_tool(name, arguments)
    
    async def list_tools(self):
        """获取可用工具列表"""
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect、connect_script或connect_sse方法")
            
        return await self.client.list_tools()
    
    async def execute_prompt(self, name: str, arguments: Optional[Dict[str, Any]] = None):
        """执行提示模板
        
        Args:
            name: 提示名称
            arguments: 提示参数
            
        Returns:
            提示执行结果
        """
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect或connect_script方法")
            
        return await self.client.execute_prompt(name, arguments)
    
    # 代理其他MCPClient的方法
    async def list_resources(self):
        """列出可用资源"""
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect或connect_script方法")
            
        return await self.client.list_resources()
    
    async def read_resource(self, uri: str):
        """读取资源内容"""
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect或connect_script方法")
            
        return await self.client.read_resource(uri)
    
    async def list_prompts(self):
        """列出可用提示模板"""
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect或connect_script方法")
            
        return await self.client.list_prompts()
    
    async def get_prompt(self, name: str, arguments: Optional[Dict[str, Any]] = None):
        """获取提示内容"""
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect或connect_script方法")
            
        return await self.client.get_prompt(name, arguments)
    
    async def ping(self):
        """Ping服务器"""
        if not self._initialized:
            await self.initialize()
            
        if not self.is_connected:
            raise RuntimeError("MCP客户端未连接，请先调用connect或connect_script方法")
            
        return await self.client.ping() 